#include <algorithm>
#include <cctype>
#include <cerrno>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <sys/stat.h>
#include <sys/types.h>
#include <vector>

#include "acl/acl.h"
#include "aclnn_segment_min_coo.h"

namespace {

constexpr int SUCCESS = 0;
constexpr int FAILED = 1;

#define INFO_LOG(fmt, args...) fprintf(stdout, "[INFO]  " fmt "\n", ##args)
#define ERROR_LOG(fmt, args...) fprintf(stderr, "[ERROR]  " fmt "\n", ##args)

struct CaseConfig {
    std::vector<int64_t> srcShape;
    std::vector<int64_t> indexShape;
    std::vector<int64_t> outShape;
    int32_t nSegments = -1;
    std::string dtype;
};

int64_t GetShapeSize(const std::vector<int64_t> &shape)
{
    int64_t total = 1;
    for (auto dim : shape) {
        total *= dim;
    }
    return total;
}

template <typename T>
bool LoadBinary(const std::string &path, std::vector<T> &data)
{
    std::ifstream file(path, std::ios::binary);
    if (!file.is_open()) {
        ERROR_LOG("无法打开文件 %s", path.c_str());
        return false;
    }
    file.seekg(0, std::ios::end);
    auto fileSize = file.tellg();
    if (fileSize < 0 || fileSize % sizeof(T) != 0) {
        ERROR_LOG("文件大小异常 %s", path.c_str());
        return false;
    }
    size_t count = static_cast<size_t>(fileSize) / sizeof(T);
    data.resize(count);
    file.seekg(0, std::ios::beg);
    file.read(reinterpret_cast<char *>(data.data()), fileSize);
    if (!file) {
        ERROR_LOG("读取失败 %s", path.c_str());
        return false;
    }
    return true;
}

template <typename T>
bool SaveBinary(const std::string &path, const std::vector<T> &data)
{
    std::ofstream file(path, std::ios::binary);
    if (!file.is_open()) {
        ERROR_LOG("无法写入文件 %s", path.c_str());
        return false;
    }
    file.write(reinterpret_cast<const char *>(data.data()), data.size() * sizeof(T));
    if (!file) {
        ERROR_LOG("写入失败 %s", path.c_str());
        return false;
    }
    return true;
}

bool EnsureDirectory(const std::string &path)
{
    struct stat st {};
    if (stat(path.c_str(), &st) == 0) {
        return S_ISDIR(st.st_mode);
    }
    if (mkdir(path.c_str(), 0755) == 0) {
        return true;
    }
    if (errno == EEXIST) {
        return true;
    }
    ERROR_LOG("创建目录失败 %s", path.c_str());
    return false;
}

int Init(int32_t deviceId, aclrtStream *stream, aclrtContext *context)
{
    auto ret = aclInit(nullptr);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("aclInit 失败, ret=%d", ret);
        return FAILED;
    }
    ret = aclrtSetDevice(deviceId);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("aclrtSetDevice 失败, ret=%d", ret);
        return FAILED;
    }
    ret = aclrtCreateContext(context, deviceId);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("aclrtCreateContext 失败, ret=%d", ret);
        return FAILED;
    }
    ret = aclrtCreateStream(stream);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("aclrtCreateStream 失败, ret=%d", ret);
        return FAILED;
    }
    return SUCCESS;
}

void Cleanup(int32_t deviceId, aclrtStream stream, aclrtContext context)
{
    if (stream != nullptr) {
        (void)aclrtDestroyStream(stream);
    }
    if (context != nullptr) {
        (void)aclrtDestroyContext(context);
    }
    (void)aclrtResetDevice(deviceId);
    (void)aclFinalize();
}

template <typename T>
int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape,
                    void **deviceAddr, aclDataType dataType, aclTensor **tensor)
{
    size_t bytes = static_cast<size_t>(GetShapeSize(shape)) * sizeof(T);
    auto ret = aclrtMalloc(deviceAddr, bytes, ACL_MEM_MALLOC_HUGE_FIRST);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("aclrtMalloc 失败, ret=%d", ret);
        return FAILED;
    }
    ret = aclrtMemcpy(*deviceAddr, bytes, hostData.data(), bytes, ACL_MEMCPY_HOST_TO_DEVICE);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("aclrtMemcpy H2D 失败, ret=%d", ret);
        return FAILED;
    }
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, nullptr, 0,
                              aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr);
    if (*tensor == nullptr) {
        ERROR_LOG("aclCreateTensor 失败");
        return FAILED;
    }
    return SUCCESS;
}

std::vector<int64_t> ParseShapeLine(const std::string &line)
{
    auto pos = line.find(':');
    if (pos == std::string::npos) {
        return {};
    }
    std::string values = line.substr(pos + 1);
    for (char &ch : values) {
        if (ch == '(' || ch == ')' || ch == '[' || ch == ']' || ch == ',') {
            ch = ' ';
        }
    }
    std::istringstream iss(values);
    std::vector<int64_t> shape;
    int64_t dim = 0;
    while (iss >> dim) {
        shape.push_back(dim);
    }
    return shape;
}

bool ParseShapeFile(const std::string &path, CaseConfig &config)
{
    std::ifstream file(path);
    if (!file.is_open()) {
        ERROR_LOG("无法打开 shape 文件 %s", path.c_str());
        return false;
    }
    std::string line;
    while (std::getline(file, line)) {
        if (line.rfind("src_shape", 0) == 0) {
            config.srcShape = ParseShapeLine(line);
        } else if (line.rfind("index_shape", 0) == 0) {
            config.indexShape = ParseShapeLine(line);
        } else if (line.rfind("out_shape", 0) == 0) {
            config.outShape = ParseShapeLine(line);
        } else if (line.rfind("dtype", 0) == 0) {
            auto pos = line.find(':');
            config.dtype = (pos == std::string::npos) ? "" : line.substr(pos + 1);
            config.dtype.erase(std::remove_if(config.dtype.begin(), config.dtype.end(), ::isspace), config.dtype.end());
        } else if (line.rfind("nSegments", 0) == 0) {
            auto pos = line.find(':');
            config.nSegments = (pos == std::string::npos) ? -1 : std::stoi(line.substr(pos + 1));
        }
    }
    return true;
}

int RunCase(const std::string &caseName,
            const std::string &inputDir,
            const std::string &outputDir,
            const CaseConfig &config,
            aclrtStream stream)
{
    std::vector<float> srcHost;
    if (!LoadBinary(inputDir + "/src_" + caseName + ".bin", srcHost)) {
        return FAILED;
    }
    std::vector<int32_t> indexHost;
    if (!LoadBinary(inputDir + "/index_" + caseName + ".bin", indexHost)) {
        return FAILED;
    }

    auto outCount = static_cast<size_t>(GetShapeSize(config.outShape));
    std::vector<float> outHost(outCount, 0.0f);
    std::vector<int32_t> argOutHost(outCount, -1);

    aclTensor *srcTensor = nullptr;
    aclTensor *indexTensor = nullptr;
    aclTensor *outTensor = nullptr;
    aclTensor *argOutTensor = nullptr;
    void *srcDevice = nullptr;
    void *indexDevice = nullptr;
    void *outDevice = nullptr;
    void *argOutDevice = nullptr;
    void *workspace = nullptr;
    uint64_t workspaceSize = 0;
    aclOpExecutor *executor = nullptr;

    int status = FAILED;
    do {
        if (CreateAclTensor(srcHost, config.srcShape, &srcDevice, aclDataType::ACL_FLOAT, &srcTensor) != SUCCESS) {
            break;
        }
        if (CreateAclTensor(indexHost, config.indexShape, &indexDevice, aclDataType::ACL_INT32, &indexTensor) != SUCCESS) {
            break;
        }
        if (CreateAclTensor(outHost, config.outShape, &outDevice, aclDataType::ACL_FLOAT, &outTensor) != SUCCESS) {
            break;
        }
        if (CreateAclTensor(argOutHost, config.outShape, &argOutDevice, aclDataType::ACL_INT32, &argOutTensor) != SUCCESS) {
            break;
        }

        auto ret = aclnnSegmentMinCooGetWorkspaceSize(srcTensor, indexTensor, config.nSegments,
                                                      outTensor, argOutTensor, &workspaceSize, &executor);
        if (ret != ACL_SUCCESS) {
            ERROR_LOG("aclnnSegmentMinCooGetWorkspaceSize 失败, ret=%d", ret);
            break;
        }
        if (workspaceSize > 0) {
            ret = aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
            if (ret != ACL_SUCCESS) {
                ERROR_LOG("workspace 分配失败, ret=%d", ret);
                break;
            }
        }

        ret = aclnnSegmentMinCoo(workspace, workspaceSize, executor, stream);
        if (ret != ACL_SUCCESS) {
            ERROR_LOG("aclnnSegmentMinCoo 执行失败, ret=%d", ret);
            break;
        }
        ret = aclrtSynchronizeStream(stream);
        if (ret != ACL_SUCCESS) {
            ERROR_LOG("aclrtSynchronizeStream 失败, ret=%d", ret);
            break;
        }

        std::vector<float> result(outCount);
        ret = aclrtMemcpy(result.data(), result.size() * sizeof(float), outDevice,
                          outCount * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST);
        if (ret != ACL_SUCCESS) {
            ERROR_LOG("拷贝 out 失败, ret=%d", ret);
            break;
        }
        std::vector<int32_t> argResult(outCount);
        ret = aclrtMemcpy(argResult.data(), argResult.size() * sizeof(int32_t), argOutDevice,
                          outCount * sizeof(int32_t), ACL_MEMCPY_DEVICE_TO_HOST);
        if (ret != ACL_SUCCESS) {
            ERROR_LOG("拷贝 arg_out 失败, ret=%d", ret);
            break;
        }

        if (!SaveBinary(outputDir + "/output_y_" + caseName + ".bin", result)) {
            break;
        }
        if (!SaveBinary(outputDir + "/output_arg_out_" + caseName + ".bin", argResult)) {
            break;
        }

        INFO_LOG("用例 %s 执行完成", caseName.c_str());
        status = SUCCESS;
    } while (false);

    if (srcTensor != nullptr) aclDestroyTensor(srcTensor);
    if (indexTensor != nullptr) aclDestroyTensor(indexTensor);
    if (outTensor != nullptr) aclDestroyTensor(outTensor);
    if (argOutTensor != nullptr) aclDestroyTensor(argOutTensor);
    if (srcDevice != nullptr) aclrtFree(srcDevice);
    if (indexDevice != nullptr) aclrtFree(indexDevice);
    if (outDevice != nullptr) aclrtFree(outDevice);
    if (argOutDevice != nullptr) aclrtFree(argOutDevice);
    if (workspace != nullptr) aclrtFree(workspace);

    return status;
}

} // namespace

int main(int argc, char **argv)
{
    int32_t deviceId = 0;
    if (argc > 1) {
        deviceId = std::atoi(argv[1]);
    }
    std::string caseName = "single_case";
    if (argc > 2) {
        caseName = argv[2];
    }
    std::string inputDir = "./input";
    if (argc > 3) {
        inputDir = argv[3];
    }
    std::string outputDir = "./output";
    if (argc > 4) {
        outputDir = argv[4];
    }

    if (!EnsureDirectory(outputDir)) {
        return FAILED;
    }

    aclrtStream stream = nullptr;
    aclrtContext context = nullptr;
    if (Init(deviceId, &stream, &context) != SUCCESS) {
        Cleanup(deviceId, stream, context);
        return FAILED;
    }

    CaseConfig config;
    if (!ParseShapeFile(inputDir + "/shape_" + caseName + ".txt", config)) {
        Cleanup(deviceId, stream, context);
        return FAILED;
    }
    if (config.dtype != "float32") {
        ERROR_LOG("当前示例仅支持 float32, shape 文件描述为 %s", config.dtype.c_str());
        Cleanup(deviceId, stream, context);
        return FAILED;
    }
    if (config.nSegments <= 0) {
        ERROR_LOG("nSegments 非法: %d", config.nSegments);
        Cleanup(deviceId, stream, context);
        return FAILED;
    }

    int status = RunCase(caseName, inputDir, outputDir, config, stream);
    Cleanup(deviceId, stream, context);
    return status;
}

